package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"main/dget"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	"github.com/sirupsen/logrus"
)

// 从 daemon.json 文件中获取所有镜像源
func getDockerRegistryMirrors() []string {
	// 默认镜像仓库
	const defaultRegistry = "https://registry-1.docker.io"
	// 获取程序的实际路径
	execPath, err := os.Executable()
	if err != nil {
		logrus.Warn("无法获取程序路径，使用默认镜像仓库: ", err)
		return []string{defaultRegistry}
	}

	// 构造 daemon.json 的路径
	daemonConfigPath := filepath.Join(filepath.Dir(execPath), "daemon.json")

	if _, err := os.Stat(daemonConfigPath); os.IsNotExist(err) {
		logrus.Warn("未找到 daemon.json 文件，使用默认镜像仓库")
		return []string{defaultRegistry}
	}

	data, err := ioutil.ReadFile(daemonConfigPath)
	if err != nil {
		logrus.Warn("读取 daemon.json 文件失败，使用默认镜像仓库: ", err)
		return []string{defaultRegistry}
	}

	var config struct {
		RegistryMirrors []string `json:"registry-mirrors"`
	}
	if err := json.Unmarshal(data, &config); err != nil {
		logrus.Warn("解析 daemon.json 文件失败，使用默认镜像仓库: ", err)
		return []string{defaultRegistry}
	}

	if len(config.RegistryMirrors) > 0 {
		return config.RegistryMirrors
	}

	return []string{defaultRegistry}
}

func main() {
	debug := flag.Bool("debug", false, "打印调试信息")
	printInfo := flag.Bool("print", false, "只打印获取信息")
	arch := flag.String("arch", "linux/amd64", "指定架构")
	proxy := flag.String("proxy", "", "http proxy")
	username := flag.String("username", "", "username")
	password := flag.String("password", "", "password")
	tags := flag.Bool("tags", false, "获取tag列表")
	syncCount := flag.Int("sync", 3, "并发下载数量")
	var registry string
	flag.StringVar(&registry, "registry", "", "指定镜像仓库")

	flag.Parse()

	if *debug {
		dget.SetLogLevel(logrus.DebugLevel)
	}

	// 获取未被 flag 包解析的参数
	args := flag.Args()

	// 校验非标志参数（包名）
	if len(args) == 0 {
		logrus.Fatalln("请输入需要下载的包名")
	}
	var pkg = args[0]

	// 解析包名和标签
	var tag string
	if strings.Contains(pkg, ":") {
		pkg, tag, _ = strings.Cut(pkg, ":")
	} else {
		tag = "latest"
	}

	// 如果包名包含仓库地址，提取仓库地址
	partsOfPkg := strings.Split(pkg, "/")
	if len(partsOfPkg) >= 3 {
		registry = "https://" + partsOfPkg[0]
		pkg = strings.Join(partsOfPkg[1:], "/")
	}

	// 如果未指定 registry，则从 daemon.json 中读取所有镜像源
	var registries []string
	if registry == "" {
		registries = getDockerRegistryMirrors()
	} else {
		registries = []string{registry}
	}

	var client dget.Client
	if *proxy != "" {
		proxyUrl, err := url.Parse(*proxy)
		if err != nil {
			logrus.Fatalln("代理地址"+*proxy+"错误", err)
		}
		logrus.Info("use http proxy ", proxyUrl.String())
		client.SetClient(&http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyURL(proxyUrl),
			},
		})
	} else {
		client.SetClient(http.DefaultClient)
	}

	// 尝试从所有镜像源下载，直到成功
	var err error
	for _, registry := range registries {
		fmt.Printf("尝试使用镜像仓库: %s\n", registry)
		err = client.Install(*syncCount, registry, pkg, tag, *arch, *printInfo, *tags, *username, *password)
		if err == nil {
			fmt.Printf("成功从镜像仓库 %s 下载\n", registry)
			break
		} else {
			logrus.Warnf("从镜像仓库 %s 下载失败: %v", registry, err)
		}
	}

	if err != nil {
		logrus.Fatalln("所有镜像源均下载失败: ", err)
	}

	os.Exit(0)
}
